"""
Scheduler implementation for message routing and agent management
"""
import json

from typing import Dict, List, Optional, Set, Callable
from ..agents.base import BaseAgent
from .message import Message, MessageType
from .routing import (MessageRouter, ConditionalRouter, BroadcastRouter,
                      RoutingConfig, RoutingVisualizer)
from .prompt import wrap_content, unwrap_message, system_prompt


class Scheduler:
    """调度器类，负责消息路由和Agent管理"""

    def __init__(self):
        self.agents: Dict[str, BaseAgent] = {}
        self._message_history: List[Message] = []
        self.max_history_size = 1000  # 最大历史消息数量

        # 初始化路由系统
        self.message_router = MessageRouter()
        self.conditional_router = ConditionalRouter()
        self.broadcast_router = BroadcastRouter()

        self.callback = None

    def register_agent(self, agent: BaseAgent) -> None:
        """
        注册一个新的Agent
        Args:
            agent: 要注册的Agent实例
        """
        if agent.agent_id in self.agents:
            raise ValueError(f"Agent ID '{agent.agent_id}' already exists")
        self.agents[agent.agent_id] = agent
        self.message_router.add_agent(agent.agent_id)  # 添加到路由系统

    def unregister_agent(self, agent_id: str) -> None:
        """
        注销一个Agent
        Args:
            agent_id: 要注销的Agent ID
        """
        if agent_id in self.agents:
            del self.agents[agent_id]

    def get_agent(self, agent_id: str) -> BaseAgent:
        """
        获取指定ID的Agent
        Args:
            agent_id: Agent ID
        Returns:
            对应的Agent实例
        """
        if agent_id not in self.agents:
            raise ValueError(f"Agent ID '{agent_id}' not found")
        return self.agents[agent_id]

    def get_all_agents(self) -> Set[str]:
        """
        获取所有已注册Agent的ID
        Returns:
            Agent ID集合
        """
        return set(self.agents.keys())

    def _protocol_convert(self, message: Message) -> Message:
        """
        将消息转换为协议格式
        Args:
            message: 要转换的消息
        Returns:
            转换后的消息字符串
        """
        msg = Message(message.content, message.sender, message.type,
                      message.receiver, message.metadata)
        if (message.sender == 'user' or not self.get_agent(
                message.sender).use_protocol) and self.get_agent(
                    message.receiver).use_protocol:
            msg.content = wrap_content(message.content,
                                       message.sender,
                                       message.receiver,
                                       message_type='常规消息')
        elif self.get_agent(message.sender).use_protocol and (
                message.receiver == 'user'
                or not self.get_agent(message.receiver).use_protocol):
            msg.content = unwrap_message(message.content)['content']
        return msg

    def route_message(self, message: Message) -> Optional[Message]:
        """
        路由单个消息到指定的接收者
        Args:
            message: 要路由的消息
        Returns:
            接收者的响应消息（如果有）
        """
        receiver_id, broadcast_targets = None, None
        # 添加到历史记录
        self._add_to_history(message)

        if message.sender != 'user' and self.get_agent(
                message.sender).use_protocol:  # 发送方使用协议
            msg = unwrap_message(message.content)
            to = msg['to'].split(' ')
            if len(to) == 1:  # 单播
                message.receiver = to[0]
                receiver_id = to[0]
                message = self._protocol_convert(message)
            elif len(to) > 1:  # 广播
                broadcast_targets = to
                message.receiver = to[0]
                receiver_id = to[0]
            else:
                raise ValueError("Invalid message format: 'to' field is empty")

        # 优先获取条件路由目标
        if not receiver_id:
            receiver_id = self.conditional_router.get_conditional_target(
                message)
        if not receiver_id:
            # receiver_id = conditional_target
            receiver_id = self.message_router.get_next_receiver(message)

        # 处理广播
        if message.sender == 'user' or not self.get_agent(
                message.sender).use_protocol:  # 发送方不使用协议
            broadcast_targets = self.broadcast_router.get_broadcast_targets(
                message)
        if broadcast_targets:
            self._handle_broadcast(message, broadcast_targets)

        # 如果没有接收者，返回None
        if not receiver_id:
            return None

        # 检查接收者是否存在
        if receiver_id not in self.agents:
            raise ValueError(f"Receiver '{receiver_id}' not found")

        # 发送消息给接收者
        receiver = self.agents[receiver_id]
        response = receiver.forward(message)

        # 如果有响应，添加到历史记录
        if response:
            self._add_to_history(response)

        return response

    def _handle_broadcast(self,
                          message: Message,
                          targets: List[str] = None) -> None:
        """
        处理广播消息
        Args:
            message: 要广播的消息
            targets: 广播目标列表
        """
        original_receiver = message.receiver
        if targets is None:
            targets = self.agents.keys()
        if message.sender in targets:
            targets.remove(message.sender)  # 移除发送者自己
        if message.receiver and message.receiver in targets:
            targets.remove(message.receiver)  # 移除消息的直接接收者

        for target_id in targets:
            if target_id in self.agents and target_id != message.sender:
                message.receiver = target_id
                msg = self._protocol_convert(message)
                self.agents[target_id].add_context(
                    json.dumps({
                        'role': 'user',
                        'content': msg.content
                    }))

        message.receiver = original_receiver

    def _add_to_history(self, message: Message) -> None:
        """
        添加消息到历史记录
        Args:
            message: 要添加的消息
        """
        self._message_history.append(message)
        if len(self._message_history) > self.max_history_size:
            self._message_history.pop(0)

    def get_message_history(self) -> List[Message]:
        """
        获取消息历史记录
        Returns:
            消息历史列表
        """
        return self._message_history.copy()

    def clear_history(self) -> None:
        """清除消息历史记录"""
        self._message_history.clear()

    def set_route(self, source: str, target: str) -> None:
        """设置点对点路由规则"""
        self.message_router.set_route(source, target)

    def set_agent_order(self, order: List[str]) -> None:
        """设置代理处理顺序"""
        self.message_router.set_default_order(order)

    def set_callback(self, callback: Callable[[Message], None]) -> None:
        """设置消息处理回调函数"""
        self.callback = callback

    def add_conditional_route(self,
                              source: str,
                              condition: Callable[[Message], bool],
                              target: str,
                              description: str = "") -> None:
        """添加条件路由规则"""
        self.conditional_router.add_condition(source, condition, target,
                                              description)

    def add_broadcast_pattern(self,
                              source: str,
                              targets: List[str],
                              condition: Optional[Callable[[Message],
                                                           bool]] = None,
                              description: str = "") -> None:
        """添加广播模式"""
        if targets is None:
            targets = self.agents.keys()
        self.broadcast_router.add_broadcast_pattern(source, targets, condition,
                                                    description)

    def export_config(self) -> RoutingConfig:
        """导出路由配置"""
        config = RoutingConfig()
        config.set_base_routes(self.message_router)
        config.set_conditional_routes(self.conditional_router)
        config.set_broadcast_patterns(self.broadcast_router)
        return config

    def import_config(self, config: RoutingConfig) -> None:
        """导入路由配置"""
        config.apply_to_router(self.message_router, self.conditional_router,
                               self.broadcast_router)

    def visualize_routes(self,
                         include_conditions: bool = True,
                         include_broadcasts: bool = True) -> str:
        """
        生成路由规则的可视化表示
        Returns:
            Mermaid图表字符串
        """
        visualizer = RoutingVisualizer(self.message_router,
                                       self.conditional_router,
                                       self.broadcast_router)
        return visualizer.generate_mermaid(include_conditions,
                                           include_broadcasts)

    def step(self, message: Message) -> Optional[Message]:
        """
        执行一步消息处理
        这是框架的核心方法，负责：
        1. 确定消息的下一个处理者
        2. 处理消息路由（包括条件路由和广播）
        3. 返回处理结果

        Args:
            message: 要处理的消息

        Returns:
            处理后的响应消息，如果没有进一步的处理则返回None
        """
        # try:
        if self.callback:
            self.callback(message)  # 调用回调函数

        # 处理消息路由
        result = self.route_message(message)

        return result

        # except Exception as e:
        #     # 发生错误时，创建系统错误消息并返回给原始发送者
        #     error_msg = Message(type=MessageType.SYSTEM,
        #                         content=f"消息处理时发生错误：{str(e)}",
        #                         sender="system",
        #                         receiver=message.sender)
        #     return error_msg

    def run_stream(self,
                   task: str,
                   cancel_token: str,
                   cancel_from: List[str],
                   max_step: int | None = None,
                   metadata: Dict = {}) -> Message | None:
        """
        自动执行任务流
        Args:
            task: 用户提示词指定的任务
            cancel_token: 终止词，当cancel_from中的任何一个代理的响应中出现此词时，任务将被终止
            cancel_from: 可以终止任务的源代理ID列表
            max_step: 最大步数，为None表示无限制
            metadata: 可选的元数据
        """
        if not self.agents:
            raise ValueError("No agents registered in the scheduler")

        for agent in self.agents.values():
            if agent.use_protocol:
                agent.sonetto_prompt = system_prompt.format(
                    members='\n'.join([
                        f'id: {agt.agent_id}, 职责: {agt.duty}'
                        for agt in self.agents.values() if agt.duty != ''
                    ]),
                    role=agent.agent_id)
                # print(agent.sonetto_prompt)

        # 创建初始消息 - 使用user作为发送者
        initial_message = Message(type=MessageType.TEXT,
                                  content=task,
                                  sender="user",
                                  metadata=metadata)
        self.set_route("user", next(iter(self.agents.keys())))

        # 添加到历史记录
        self._add_to_history(initial_message)

        # 初始化步数计数器
        step_count = 0
        current_message = initial_message

        # 主执行循环
        while True:
            # 检查是否达到最大步数
            if max_step is not None and step_count >= max_step:
                break

            # 执行一步消息处理
            response = self.step(current_message)
            if isinstance(response, str):
                print(response.sender, ':\n' + response.content)
            else:
                print(response.sender, ':\n')
                print(response.content)
                print('to', response.receiver)
            print('-------------------------------------')
            # input('按回车继续')
            step_count += 1

            # 如果没有进一步的响应，则终止循环
            if not response:
                break

            # 检查是否包含取消令牌
            if response.sender in cancel_from:
                # 将内容转换为字符串检查取消令牌
                content_str = str(
                    response.content) if response.content is not None else ""
                if cancel_token in content_str:
                    break

            # 将响应作为下一个消息继续处理
            current_message = response

        return response
